package jwe_test

import (
	"context"
	"testing"

	"github.com/lestrrat-go/jwx/v3/internal/json"
	"github.com/lestrrat-go/jwx/v3/jwa"
	"github.com/lestrrat-go/jwx/v3/jwe"
	"github.com/stretchr/testify/require"
)

// Pin represents the structured clevis data which can be used to decrypt the jwe message
type Pin struct {
	Pin     string      `json:"pin"`
	Tang    *TangPin    `json:"tang,omitempty"`
	Tpm2    *Tpm2Pin    `json:"tpm2,omitempty"`
	Sss     *SssPin     `json:"sss,omitempty"`
	Yubikey *YubikeyPin `json:"yubikey,omitempty"`
}

type TangPin struct {
	Advertisement *json.RawMessage `json:"adv,omitempty"`
	URL           string           `json:"url"`
}

type Tpm2Pin struct {
	Hash    string `json:"hash,omitempty"`
	Key     string `json:"key,omitempty"`
	JwkPub  string `json:"jwk_pub,omitempty"`
	JwkPriv string `json:"jwk_priv,omitempty"`
	PcrBank string `json:"pcr_bank,omitempty"`
	PcrIds  string `json:"pcr_ids,omitempty"`
}

type SssPin struct {
	Jwe       []string `json:"jwe"`
	Threshold int      `json:"t"`
	Prime     string   `json:"p"`
}

type YubikeyPin struct {
	Type      string     `json:"type"`
	Challenge string     `json:"chalelenge"`
	Slot      int        `json:"slot"`
	Kdf       YubikeyKdf `json:"kdf"`
}

type YubikeyKdf struct {
	Type       string `json:"type"`
	Hash       string `json:"hash"`
	Iterations int    `json:"iter"`
	Salt       string `json:"salt"`
}

func TestGH402(t *testing.T) {
	key := []byte{195, 170, 42, 171, 98, 176, 98, 162, 57, 170, 62, 69, 175, 209, 200, 151, 81, 135, 63, 43, 93, 20, 16, 111, 13, 26, 138, 188, 15, 19, 26, 242}
	data := "eyJhbGciOiJkaXIiLCJjbGV2aXMiOnsicGluIjoic3NzIiwic3NzIjp7Imp3ZSI6WyJleUpoYkdjaU9pSkZRMFJJTFVWVElpd2lZMnhsZG1seklqcDdJbkJwYmlJNkluUmhibWNpTENKMFlXNW5JanA3SW1Ga2RpSTZleUpyWlhseklqcGJleUpoYkdjaU9pSkZRMDFTSWl3aVkzSjJJam9pVUMwMU1qRWlMQ0pyWlhsZmIzQnpJanBiSW1SbGNtbDJaVXRsZVNKZExDSnJkSGtpT2lKRlF5SXNJbmdpT2lKQlZqbEplVmszZDNCamNtMVpZMDkxZWsxTE9VSjJaMDFDZDNJd09FNXhURXBtUW5Ca1IzVXpNa1ZDWjJORmRITnpNRGswVEdweFdXSlFhVmhQTFhWTWMyeDVValZUVlRsdlNHbGhaRlI1VlhWb1UzUTJhbXQzSWl3aWVTSTZJa0ZIVkRseVpYcHlkV0pITFU1eFgxWkxPRUl3UXpGTE9WSTJNV05YU1U1UFNVaEZSMU5MYlZKclUxZEtjVEJRWHpkaFVHaDJlV3hhWDNaTFdHSk5TMGhrU1d0VFFrd3dSMWRRTmxaaVdqbExZVEJHVkZGSFpqTWlmU3g3SW1Gc1p5STZJa1ZUTlRFeUlpd2lZM0oySWpvaVVDMDFNakVpTENKclpYbGZiM0J6SWpwYkluWmxjbWxtZVNKZExDSnJkSGtpT2lKRlF5SXNJbmdpT2lKQlNIZFFhR2hPZDBac1ptc3RNR2d6VDNGS1JFVXpkekpVWDJGcExUbGlZWGhXY0VKa05TMWtNbXhIU0VseFgyTXhVSEZrTW1RMlYxSkRXbEZVYkc5cWQyUmhOblJUUkVndFUxTmhaVE5aTmxGVlZIVkVla1puSWl3aWVTSTZJa0ZrYVhkdVozaDVlamt4UVMxUFRGODFhekIwVEcwMlZtSTNSVkZxWVcxQ1NWUlhlVzVuUldSRVpuRlRSbm8xWlRkdVkycFdlR3BpWDFCSVMwbFNaRzFPT1hCa2RFMTRlRUpPYTFkV2RsOWZUbkZxVTI5SU5UUWlmVjE5TENKMWNtd2lPaUpvZEhSd09pOHZiRzlqWVd4b2IzTjBPak16TURJeEluMTlMQ0psYm1NaU9pSkJNalUyUjBOTklpd2laWEJySWpwN0ltTnlkaUk2SWxBdE5USXhJaXdpYTNSNUlqb2lSVU1pTENKNElqb2lRVWR1YjNrMFQyMTBhRWt3TjNKV1QxbFNTVGhxUTIxVGMxQjZXVVYwUkcwNWNqRkRaa0ptVmxCa2RscEdVVE0wY2taVmJFTTFRbGRrVlZweVMyOXNYMHh5UVhsdFlrY3pWVEZhY1VRMU9YSnBlVGQzV1hGWVV5SXNJbmtpT2lKQlNHaHNSVGx0T1hoR04xVkVVVXd4WjJSd2RGaE5Zbm94VUhNMloweDRObWx1ZGpodU4xVXpUamhpU0dSTlJIRldhalZKYVdGVGFGbDRZM1kyYzFSQ2VuUXpNMmxrUlRaZlYyUm1Oell3WVhkV1VVUkRXa1oxSW4wc0ltdHBaQ0k2SWpWQlNVaHJjblpDYzFOZmQzRjFMVGN5TUd0Q1dtMTBWblpxYkdsVGVWSmpUVXREVFU5dmJuWlFSVWtpZlEuLlR0SzJzWTE2bDJnSmx6c18uRFZqcGVkbGUxY29xZkFkLUZ3bnM5RmtRVkR0ak4zWU1NWjc4VjNLVmI3VWpFd2p1eTBXYjZ0eUxzTmZsVnJaVU85dDM1X2pfSlpobU0ycllfV1RyOHcuVkZMYjFvcTFqZzRsdnJpY3laaXF6USIsImV5SmhiR2NpT2lKRlEwUklMVVZUSWl3aVkyeGxkbWx6SWpwN0luQnBiaUk2SW5SaGJtY2lMQ0owWVc1bklqcDdJbUZrZGlJNmV5SnJaWGx6SWpwYmV5SmhiR2NpT2lKRlEwMVNJaXdpWTNKMklqb2lVQzAxTWpFaUxDSnJaWGxmYjNCeklqcGJJbVJsY21sMlpVdGxlU0pkTENKcmRIa2lPaUpGUXlJc0luZ2lPaUpCVUZWTVMxVllhMUp2TjJOdGFXOXFVRUV3ZFVZMGVIcFFXakJ0V21sQ1NFeFBhVWRyV0hkQ1pGTkVibVpNV1ZaTmJFeFZWR0UzTXkxNWJYSnVWME5mUlRZdE9DMWpUbnA0U1RaMWVtMTRhVFphTXpCUU1EaHhJaXdpZVNJNklrRlBPVTFWTVdGQmRWRnJNVEJrY1dGbVIwRlFYMlpKVGkwelRXTldNRkpGT0drNGFVdG5UbGhuWm14SlFtbDZNRGMxVmpCWVJIZEZNbkpSU2xka1VsRktlWGczT0hFeVJrTndkekZGTTFCMk1YRnZZMmhCZDJnaWZTeDdJbUZzWnlJNklrVlROVEV5SWl3aVkzSjJJam9pVUMwMU1qRWlMQ0pyWlhsZmIzQnpJanBiSW5abGNtbG1lU0pkTENKcmRIa2lPaUpGUXlJc0luZ2lPaUpCVUUxSFZYWTFjSFZUZWt0MU1HRjRZVFJRYldOeFRsQkxiMjR0TW1oYVVGSjJXbmRwV2xkc05sOVNNREJ2YlZsMldESXplRlJ4UldGUVdFMXphemxFYm1SdGNrUnpWbEZvTlRFMFEzTTBSbGszV0VaTlkyUkVJaXdpZVNJNklrRlZNa051VUV0RlJVeDBhSFZaYlZaeWJWZEJORFZLYjNocE1FRkhjM2w1T0daak5UTlJiMHA2UjA4NWNYVnRXVVpRTTFWbVoyWjZibkZsTnpWcE5uUXpkSFJPUW0xSWNGVlhTa1JuUkRkdGRGbEdVbXRzYzFjaWZWMTlMQ0oxY213aU9pSm9kSFJ3T2k4dmJHOWpZV3hvYjNOME9qUXhPRFl4SW4xOUxDSmxibU1pT2lKQk1qVTJSME5OSWl3aVpYQnJJanA3SW1OeWRpSTZJbEF0TlRJeElpd2lhM1I1SWpvaVJVTWlMQ0o0SWpvaVFWazJVMmxoUzFaRFMwUnFRMjUwTFdWNFJFZzJTa3RhYVdkaWRXOVpZVzV3Y0hWRVkxOXNOVXBxY1ROR05GWnpMV3hJYzFGbVYzQmFOVXRPWTNCUWFITmZObFpVUmt4ZlYwMVdja2xCZG5wWWVFeExZMmhKY0NJc0lua2lPaUpCU1RCVFZVSTFUMG80ZGtkR05YRndXRzVFZVRkd1pteFBNVjlsZW14WWQyTTNValpxY25RelJrSklPRTEwZG0xcFRHVkZNRzR0VEhZNFMyNWZMVWMzUlhJM1MwRklNSEoyWW5kMlVFZzBjbXhrU1VrM2RETmxJbjBzSW10cFpDSTZJa1JNT1ZaclRHcFVaVXBKTVhRd2ExUnNUbnBzVGxrM1RHcDFTWGxLZHpGUE1HTkdYMHN0YW5GdGFHOGlmUS4uSjJ4c2VjeTN3c2FGRjFfdi5uQzdPSzVrdXAwbkgxVGwwOFBRUE82QTllNjFYMXNGY1RFUVZjTjNkd0NWYTFMZ2p2STFfSWZqTU5TX2cweTI4WGM1ZThQZ2hFS01UQ0hkQk44UzFBdy5yV0s5Qi0yT2RtWUlIbmQzMzcxeHpRIiwiZXlKaGJHY2lPaUpGUTBSSUxVVlRJaXdpWTJ4bGRtbHpJanA3SW5CcGJpSTZJblJoYm1jaUxDSjBZVzVuSWpwN0ltRmtkaUk2ZXlKclpYbHpJanBiZXlKaGJHY2lPaUpGUTAxU0lpd2lZM0oySWpvaVVDMDFNakVpTENKclpYbGZiM0J6SWpwYkltUmxjbWwyWlV0bGVTSmRMQ0pyZEhraU9pSkZReUlzSW5naU9pSkJRbVZ1WVRaWVVUQkNTMmQwU1U5c05FMXZUakZEZFhKc2REVk1lVEZ6YUZNMFRFbHdkV05aUkVaWE9VNDVNRWgxY1RSUk4wcENPR3RNVDNabVdEUnBhSFJwTFdWQ1JXdGpVR3BCYzNWWU9HWm5kbVEyVEZSM0lpd2llU0k2SWtGUlRFWnVZMmhmZVRCaVdFNTRTRGt6UW10NWQwSmFhbXB6T1ROaFptOXFXWGhyV0d0dFFXZFRVbnA1ZGxaWGFrbHJRMEYwVWtoRVZuVmhlbkUxTWtSRFYxVnRlRU5YWm1OR05qWkdSaTB5Y1ZrellsTk9Na01pZlN4N0ltRnNaeUk2SWtWVE5URXlJaXdpWTNKMklqb2lVQzAxTWpFaUxDSnJaWGxmYjNCeklqcGJJblpsY21sbWVTSmRMQ0pyZEhraU9pSkZReUlzSW5naU9pSkJaR2hHZW05RE5FZG5WRk5CWlUxTFVYSlhjRWRCTWtndFEyeEVNVE52U1d4SVprczFTWEI2ZVZseWRuVm1Sbk5DZDJkYWJWYzRWR0ZXUkdSYVltbGlkM3BtTldSMk1rOHRkVkpmU2pSS2IyNDRNekJ0ZFVGa0lpd2llU0k2SWtGUlpFMVdka3RGWW5obWMwbHRiVFJtVHpZMlFsOHdVMnhPWDBaRmVEQkthVlpyVjFKQlFrRnBhRVExV1VKUWVVbEpjM1pGUjJoR1dqaFpOVXBWTXpoQ1JWVlJUWGhyVkZsT1pEVlpOR1UxVG1STVVVbDNObDhpZlYxOUxDSjFjbXdpT2lKb2RIUndPaTh2Ykc5allXeG9iM04wT2pNME1qVTVJbjE5TENKbGJtTWlPaUpCTWpVMlIwTk5JaXdpWlhCcklqcDdJbU55ZGlJNklsQXROVEl4SWl3aWEzUjVJam9pUlVNaUxDSjRJam9pUVZaU1praEpWVVJ2U0RCb1FYRnhOSFpJYVVGMmNVOVNXblJUVFZkS2VGVmtPWEozV2tWdVZFMUVWbWt4VG5KbGRWTnZWWHBNWmxkcFVVcDBiR3BxU21ZM1NqRnpUbGRVUmpKaVFtTnJNSEJrUnpkcFJHVnhVeUlzSW5raU9pSkJSak5aUms5dlNXNW1jVU5yVFRKa1pISnlkV1JyY21zMmN6Wk1SMUpNU0U5YVl6QlhhMmRhZFZsb1FsRk9ZbVZXY25sWFZTMW9aMTgwYkhCRE5EUjRaRmh3ZUc4NWFGOWhSVlZ0U0V0UGFWOWpVM2xhY1VjM0luMHNJbXRwWkNJNkluQnNablF4UW5WR2EyeGZaRGhyZFZOaWNVVlVSbXhxYjJwM01WZDBMVFpLYTJ0VGVXOWtWa1JEZDNjaWZRLi5CUUVtRXFhc1FHNHNKY1d2Lmo5RzBweVEwbTY0TU5OTEZicm4wT01BbmZaOTJ5bWNxV19ZM2ZIMWJEd2JFbnI1U1FHWF9jWjNfbUpGSS0zMFVJOEZuZzRUaXRxMU9JT2JySktETTh3LndRTmVJVzhzZnNua1V1aTUwTWpYd0EiLCJleUpoYkdjaU9pSkZRMFJJTFVWVElpd2lZMnhsZG1seklqcDdJbkJwYmlJNkluUmhibWNpTENKMFlXNW5JanA3SW1Ga2RpSTZleUpyWlhseklqcGJleUpoYkdjaU9pSkZRMDFTSWl3aVkzSjJJam9pVUMwMU1qRWlMQ0pyWlhsZmIzQnpJanBiSW1SbGNtbDJaVXRsZVNKZExDSnJkSGtpT2lKRlF5SXNJbmdpT2lKQlJVdzRjbkZ3UjE4dE9FRXRiM1UwYTJaMFJXbFNjbVZ5VldOb1ozWmlVRzFXYUV0bVJVNWtXVXd0ZDNOVlRucGZNMlpRZW5GUVVtWkhiRzltTkZkR2RFMTNNM0I1YnpaSVFUaFFTM0JHVlc5VWJ6ZHRjbWRSSWl3aWVTSTZJa0ZDUWxsVU9WOWpWRnBuT1RCWlJXWk9jWFV4V2tkSGFtaFphMWxwTTFkWlRqWmpialZJUmtjdGJ5MUdOR3BIWWpkVFpEazFTbWsxTkhkVWRUQmxOVmM1UjE5cGVFRlZYMWh0Wm1GSmQyZFRXVXN3WDJScVpEVWlmU3g3SW1Gc1p5STZJa1ZUTlRFeUlpd2lZM0oySWpvaVVDMDFNakVpTENKclpYbGZiM0J6SWpwYkluWmxjbWxtZVNKZExDSnJkSGtpT2lKRlF5SXNJbmdpT2lKQlZYSktNME5GUzFZd2FGWlNNamhLWTE5eWRXUkpUWEF0YmxCUmRtdFNSemxGVlZsVVRrTjJiVEZrY2tWUlJVNXVWMjB3VXpNelMycEtTMWx1VFhObVJrTlRPR3BYV0hJMFkxUTNTMFV4V2xkWmFGRnJkRzF3SWl3aWVTSTZJa0ZrYUVkWGFFVjJNREE1VjJwdFRUQjFZWFpFTm5sR01HczNMVXR0WkRCTmNHNTNRelZtWVVReFZuTkhWSGhLVG5SYVptSmlWbGhGVTIxeVJYZ3pYMVJoVERZMFpVeFhObFZwTFVveWJscHZhbEpUU0hWWmFtTWlmVjE5TENKMWNtd2lPaUpvZEhSd09pOHZiRzlqWVd4b2IzTjBPalF6TVRVMUluMTlMQ0psYm1NaU9pSkJNalUyUjBOTklpd2laWEJySWpwN0ltTnlkaUk2SWxBdE5USXhJaXdpYTNSNUlqb2lSVU1pTENKNElqb2lRV1JrT0Rnd05FdFlRbGhuWmpacU1EQkdVbEIyUzBsWmFGaE5ZMnRpWjE5c1NXbHdhbVF6TlhGUk9FVk9MVEZGUTNOUVZrNHRTa1I1ZW1sVmVuSkhjVTVOUkRGSmJGOUlTMUIyVGpaUmFVMU9kVXB0VVcxcWJpSXNJbmtpT2lKQlF5MW9Tak5DUTBvMU9VSnZhbWR1Ym5oZlZuaHBWekZNTVU5UlYxTXRZMmxzV2xkRGVIaGhZVUpyYjA0MVV6RnpiMnh4Y1ZwbVFuQkNiakZSWW5OTlJEZGZiVFEwVm1weGRWQlFjMjQwVmpFNFFqQlZkV2hDSW4wc0ltdHBaQ0k2SWtaUmNWZDBUWGQ0V2pWbVEyWjFSekp1WldWT1lreEpkQzF6Tm1Ka2RUQXpaMUZSU0Vvd00wY3RNMDBpZlEuLmduekZoT21WQk41UFNWT2kuZlBkQWJ2YlgtLUFXV1pPREZEdy1Yd2VvZHJQamc5UGM4SlFQQjIwZTlPQ2Y5RTZmSEtNYUJlR2xVSWRyOHViRjdRUjllM3Vnalk2ZEdLOEoyeklyYXcuTlV6T3M4a01HVVhMa3BXNWxja3BnZyIsImV5SmhiR2NpT2lKRlEwUklMVVZUSWl3aVkyeGxkbWx6SWpwN0luQnBiaUk2SW5SaGJtY2lMQ0owWVc1bklqcDdJbUZrZGlJNmV5SnJaWGx6SWpwYmV5SmhiR2NpT2lKRlEwMVNJaXdpWTNKMklqb2lVQzAxTWpFaUxDSnJaWGxmYjNCeklqcGJJbVJsY21sMlpVdGxlU0pkTENKcmRIa2lPaUpGUXlJc0luZ2lPaUpCWTNVeWFFOTVXSGRPVG1OYVpXRnJkVTFRVW01aGFqSnJXbGxJUldkMVFuRmZiR0ZMU0VZeFUwbzFjMVp3UmswMVVsZFJYMVp5U2pKVmVtMHpXa0k0T1Y5Nk5HZExVVzVETFU1TFJHTk5NMHcxUkVsdlh6UlNJaXdpZVNJNklrRmlOVXBTU1hOeU5HSm9kRXhvV1dabVVITlpaMUZqT0hocGNXcFplak5wUm1GeVpIbEVkVFZRZHpSVFRIZExWR1JuV2tSMGNtODVNSHByTVRKVFpERmhYemgxZVRkV1drMHRTRVkzUWs5bGNXeG9RV3R2VEdnaWZTeDdJbUZzWnlJNklrVlROVEV5SWl3aVkzSjJJam9pVUMwMU1qRWlMQ0pyWlhsZmIzQnpJanBiSW5abGNtbG1lU0pkTENKcmRIa2lPaUpGUXlJc0luZ2lPaUpCVm1Sa1ZFWklNRWxwTm1OWk1sQk9XVVp1WVZkT2IwaENlRzVuVFZaR1kzZFlObVYzU25vd1VVUlZjbkF5UmxsT1pVNHpaQzFxYWxwcVdYQmZkRkJ0Ukdjd05YcG1iSHAyU0VoSWJscFlUM2RLT1RSSU5VOHhJaXdpZVNJNklrRllaRkZ5VDBFdFJXWjFTM2xHU0drdE5sVlFSRGMwUnkxaVdVdFhNbFV4YVZselpFTkRSakp5WWs5T1JXeDBTVFowT0ZSZk56bHZjR1pLTkcxU1gySnliMVZqVVhBNU1GOWFZM05RY3poZlZtRndSbmhyVmpJaWZWMTlMQ0oxY213aU9pSm9kSFJ3T2k4dmJHOWpZV3hvYjNOME9qTTBPRGM1SW4xOUxDSmxibU1pT2lKQk1qVTJSME5OSWl3aVpYQnJJanA3SW1OeWRpSTZJbEF0TlRJeElpd2lhM1I1SWpvaVJVTWlMQ0o0SWpvaVFXUkNVVlJKUkdjNGNEWm5NVzh4T0d4R1dYRm1kMXBhTWxaaFJteG9VMlF0YVhkV09EQnRXalZaWm5wWU5qbFNWV3BxYzNKc1ZHZEtjbXMwY2pKaWRrSlZUbE5YY0cxbVEyeHBYek0xYm1Wb2NrZHpRbW94VUNJc0lua2lPaUpCVTJGVVYxY3RSRk5ZVVhObFgxUmxibDg0VkdkVmFrWTNVRWxRTm5wcVpUUmxSbk55YkY5WFRESkhNRGwzTFROalYzVlRlVlZmUzFvemF6VndTVlp0TlZkaU1tcGlXREptV2sxWmVFNVpRbDh5Vm10d1QxWnVJbjBzSW10cFpDSTZJbEk0YWtWSk1YQlRNVWxoY0hWVVZHTktObGhJWkROWFUzbGZRbmhXU1hWb1VVeDFVekJwYjJkU09UUWlmUS4uTEk2ZC1TNUlYSDRIY2lxeS5ESWl3M1hFeXlaMVZWZDFPcWYwS1pJeGZrUWtnZXo4bXl2N0dfZGFMOWxIdjVFVlFMQzREVG9PMWF6UDQ5SDNGUk5PRHZIcXBmaWp1a3VDX2JmLVZldy41NGpyZnZWdEYxR3ItYU9TVEFjSWV3Il0sInAiOiI0NERteG1EZVFiWTJ3eHpIYm9PODFkbTBCbENNcVJzeEtLQ0NDQWo4TVVjIiwidCI6M319LCJlbmMiOiJBMjU2R0NNIn0..zz3fUXsiaME2cSoy.LTQovHUvDP4MXT2_sHgf_cM2gicobD5kGXEl5eY.MK3Lf6IwaoVUvCTp1Q5VOA"

	decrypt := func(customField bool) {
		t.Helper()
		m := jwe.NewMessage()
		// Test WithPostParse while we're at it
		plain, err := jwe.Decrypt([]byte(data),
			// This is a really cheesy way of creating a jwa.KeyEncryptionAlgorithm
			// but a bogus one.
			jwe.WithKey(jwa.NewKeyEncryptionAlgorithm("invalid algorithm"), nil),
			jwe.WithMessage(m),
			jwe.WithKeyProvider(jwe.KeyProviderFunc(func(_ context.Context, sink jwe.KeySink, _ jwe.Recipient, _ *jwe.Message) error {
				sink.Key(jwa.DIRECT(), key)
				return nil
			})),
		)

		require.NoError(t, err, `jwe.Decrypt should succeed`)

		if string(plain) != "testing Shamir Secret Sharing" {
			t.Errorf("expected 'testing Shamir Secret Sharing', got %s", string(plain))
			return
		}

		if customField {
			require.NotNil(t, m.ProtectedHeaders(), `m.ProtectedHeaders should be non-nil`)

			var v Pin
			require.NoError(t, m.ProtectedHeaders().Get("clevis", &v), `m.Get("clevis") should be succeed`)
		}
	}
	decrypt(false)

	// register field deserialized and run decryption again
	jwe.RegisterCustomField("clevis", Pin{})
	decrypt(true) // used to fail before, but this should pass
}
